#include <ros/ros.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_types.h>
#include <sensor_msgs/PointCloud2.h>
#include <pcl_conversions/pcl_conversions.h>
#include <rosbag/bag.h>
#include <rosbag/view.h>
#include <boost/foreach.hpp>
#include <sys/stat.h>
#include <livox_ros_driver2/CustomMsg.h>  // 引入自定义消息头文件

#define foreach BOOST_FOREACH

// 转换函数：将 CustomMsg 转换为 PointCloud2
void convertCustomMsgToPointCloud2(const livox_ros_driver2::CustomMsg& custom_msg, sensor_msgs::PointCloud2& cloud_out) {
    // 创建一个新的 PointCloud2 消息
    sensor_msgs::PointCloud2 cloud;
    cloud.header = custom_msg.header;
    cloud.height = 1;  // 假设点云是一维的（非结构化）
    cloud.width = custom_msg.point_num;  // 点的数量
    cloud.is_dense = false;
    cloud.is_bigendian = false;

    // 设置字段信息
    sensor_msgs::PointField field_x;
    field_x.name = "x";
    field_x.offset = 0;
    field_x.datatype = sensor_msgs::PointField::FLOAT32;
    field_x.count = 1;

    sensor_msgs::PointField field_y;
    field_y.name = "y";
    field_y.offset = 4;
    field_y.datatype = sensor_msgs::PointField::FLOAT32;
    field_y.count = 1;

    sensor_msgs::PointField field_z;
    field_z.name = "z";
    field_z.offset = 8;
    field_z.datatype = sensor_msgs::PointField::FLOAT32;
    field_z.count = 1;

    sensor_msgs::PointField field_reflectivity;  // 如果需要反射率
    field_reflectivity.name = "reflectivity";
    field_reflectivity.offset = 12;
    field_reflectivity.datatype = sensor_msgs::PointField::FLOAT32;
    field_reflectivity.count = 1;

    cloud.fields.resize(4);  // 如果不需要反射率，改为3
    cloud.fields[0] = field_x;
    cloud.fields[1] = field_y;
    cloud.fields[2] = field_z;
    cloud.fields[3] = field_reflectivity;  // 如果不需要反射率，删除这一行

    cloud.point_step = 16;  // 每个点占用 16 字节 (4 * float32)，如果不需要反射率，则为12字节
    cloud.row_step = cloud.point_step * cloud.width;

    // 分配数据缓冲区
    cloud.data.resize(cloud.row_step);

    // 填充点云数据
    for (size_t i = 0; i < custom_msg.points.size(); ++i) {
        memcpy(&cloud.data[i * cloud.point_step + 0], &custom_msg.points[i].x, sizeof(float));  // x
        memcpy(&cloud.data[i * cloud.point_step + 4], &custom_msg.points[i].y, sizeof(float));  // y
        memcpy(&cloud.data[i * cloud.point_step + 8], &custom_msg.points[i].z, sizeof(float));  // z
        if (cloud.fields.size() == 4) {  // 如果需要反射率
            memcpy(&cloud.data[i * cloud.point_step + 12], &custom_msg.points[i].reflectivity, sizeof(float));  // reflectivity
        }
    }

    cloud_out = cloud;
}

int main(int argc, char** argv)
{
    ros::init(argc, argv, "pcd_extractor");
    ROS_INFO("Program started.");

    if (argc < 3)
    {
        ROS_INFO("Usage: %s <bag_file> <output_dir>", argv[0]);
        return 1;
    }

    std::string bag_file = argv[1];
    std::string output_dir = argv[2];

    ROS_INFO("Bag file: %s", bag_file.c_str());
    ROS_INFO("Output directory: %s", output_dir.c_str());

    rosbag::Bag bag;
    try
    {
        bag.open(bag_file, rosbag::bagmode::Read);
        ROS_INFO("Successfully opened the bag file.");
    }
    catch (rosbag::BagIOException& e)
    {
        ROS_ERROR_STREAM("Failed to open bag file: " << e.what());
        return 1;
    }

    // 检查并创建输出目录
    if (access(output_dir.c_str(), F_OK) == -1) {
        if (mkdir(output_dir.c_str(), 0755) == -1) {
            ROS_ERROR_STREAM("Failed to create output directory: " << output_dir);
            return 1;
        } else {
            ROS_INFO_STREAM("Created output directory: " << output_dir);
        }
    }

    std::vector<std::string> topics{"/livox/lidar"};
    rosbag::View view(bag, rosbag::TopicQuery(topics));
    ROS_INFO("Starting to process messages.");

    int i = 0;
    foreach(rosbag::MessageInstance const m, view)
    {
        // 尝试实例化自定义消息类型
        livox_ros_driver2::CustomMsg::ConstPtr custom_msg = m.instantiate<livox_ros_driver2::CustomMsg>();
        if (custom_msg != nullptr)
        {
            ROS_INFO("Processing frame %d.", i);

            sensor_msgs::PointCloud2 cloud_msg;
            convertCustomMsgToPointCloud2(*custom_msg, cloud_msg);  // 转换自定义消息到 PointCloud2

            pcl::PCLPointCloud2 pcl_pc2;
            pcl_conversions::toPCL(cloud_msg, pcl_pc2);

            pcl::PointCloud<pcl::PointXYZ>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZ>);
            pcl::fromPCLPointCloud2(pcl_pc2, *cloud);

            ROS_INFO("Converted frame %d to PCL format with %lu points.", i, cloud->points.size());

            std::stringstream ss;
            ss << output_dir << "/frame_" << i << ".pcd";
            ROS_INFO("Attempting to save PCD file to: %s", ss.str().c_str());
            if (pcl::io::savePCDFileASCII(ss.str(), *cloud) == 0)
            {
                ROS_INFO("Successfully saved frame_%d.pcd", i);
            }
            else
            {
                ROS_ERROR("Failed to save frame_%d.pcd", i);
            }

            ++i;
        }
        else
        {
            ROS_WARN("Failed to instantiate CustomMsg message.");
        }
    }

    bag.close();
    return 0;
}